package com.jvxb.manage.livable.util;

import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLLimit;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.statement.*;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlInsertStatement;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlRenameTableStatement;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectQueryBlock;
import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlSchemaStatVisitor;
import com.alibaba.druid.sql.parser.SQLParserUtils;
import com.alibaba.druid.sql.parser.SQLStatementParser;
import com.alibaba.druid.stat.TableStat;
import com.jvxb.common.constants.DbConstants;

import java.util.*;

/**
 * @author jvxb
 * @since 2022-02-09
 */
public class SqlCheckUtil {

    public static String CHARSET = "CHARSET";
    public static String CHARACTER_SET = "CHARACTER SET";
    public static String DDL = "DDL";
    public static String DML = "DML";
    public static String USE = "USE";

    private SqlCheckUtil() {

    }

    public static String checkSqlType(List<SQLStatement> sqlStatements) {
        boolean ddl = false;
        boolean alter = false;
        boolean create = false;
        for (SQLStatement sqlStatement : sqlStatements) {
            if (!(sqlStatement instanceof SQLUseStatement)
                    && !(sqlStatement instanceof SQLCreateTableStatement)
                    && !(sqlStatement instanceof SQLAlterStatement)
                    && !(sqlStatement instanceof SQLInsertStatement)
                    && !(sqlStatement instanceof SQLDeleteStatement)
                    && !(sqlStatement instanceof SQLUpdateStatement)
                    ) {
                throw new RuntimeException("不支持的变更语句：" + sqlStatement.toString());
            }

            if (sqlStatement instanceof SQLAlterStatement) {
                alter = true;
                ddl = true;
            } else if (sqlStatement instanceof SQLCreateStatement) {
                create = true;
                ddl = true;
            } else if (sqlStatement instanceof SQLDropStatement) {
                ddl = true;
            }
        }
        if (alter && create) {
            throw new RuntimeException("Alter和Create语句不能一起提交");
        }
        if (ddl) {
            for (SQLStatement sqlStatement : sqlStatements) {
                if (create) {
                    if (sqlStatement instanceof SQLDeleteStatement || sqlStatement instanceof SQLUpdateStatement || sqlStatement instanceof SQLReplaceStatement) {
                        throw new RuntimeException("DDL语句不能和DML语句一起提交(除create和insert外)");
                    }
                } else if (sqlStatement instanceof SQLInsertStatement || sqlStatement instanceof SQLDeleteStatement || sqlStatement instanceof SQLUpdateStatement || sqlStatement instanceof SQLReplaceStatement) {
                    throw new RuntimeException("DDL语句不能和DML语句一起提交(除create和insert外)");
                }
            }
        }
        return null;
    }

    public static String checkAlterDb(List<SQLStatement> sqlStatements) {
        String lastDbName = "";
        for (SQLStatement sqlStatement : sqlStatements) {
            if (!sqlStatement.toString().endsWith(";")) {
                throw new RuntimeException("sql请以分号结尾");
            }
            //USE语句 获取库名;
            if (sqlStatement instanceof SQLUseStatement) {
                String useDB = sqlStatement.toString();
                lastDbName = useDB.substring(useDB.indexOf(' ') + 1, useDB.lastIndexOf(';'));
                continue;
            }

            //获得当前语句表名
            MySqlSchemaStatVisitor visitor = new MySqlSchemaStatVisitor();
            sqlStatement.accept(visitor);
            Map<TableStat.Name, TableStat> tables = visitor.getTables();
            Set<TableStat.Name> names = tables.keySet();
            String table = "";
            for (TableStat.Name name : names) {
                table = name.toString();
            }

            if (!table.contains(".") && lastDbName.length() == 0) {
                throw new RuntimeException("表`" + table + "`未指定数据库");
            }
        }

        for (int i = 0; i < sqlStatements.size(); i++) {
            SQLStatement sqlStatement = sqlStatements.get(i);
            if (sqlStatement instanceof SQLInsertStatement) {
                SqlCheckUtil.checkDuplicate(sqlStatement);
            } else if (sqlStatement instanceof SQLAlterStatement) {
                String table = getTableName(sqlStatements, i);
                String[] split = table.split("\\.");
                if (split.length == 2 && split[0].equalsIgnoreCase("information_schema")) {
                    throw new RuntimeException("不允许对`information_schema`库进行操作");
                }
            }
        }
        return null;
    }

    public static void checkDuplicate(SQLStatement sqlStatement) {
        MySqlInsertStatement sqlInsertStatement = (MySqlInsertStatement) sqlStatement;
        List<SQLExpr> duplicateKeyUpdate = sqlInsertStatement.getDuplicateKeyUpdate();
        if (duplicateKeyUpdate.size() > 0) {
            throw new RuntimeException("不支持使用 ON DUPLICATE KEY,语句是不安全的操作，可能引发MySQL主从数据不一致");
        }
        SQLSelect query = sqlInsertStatement.getQuery();
        if (query != null) {
            throw new RuntimeException("暂不支持 INSERT INTO SELECT 语句");
        }
    }

    public static String getTableName(List<SQLStatement> sqlStatements, int i) {
        //获得表名
        MySqlSchemaStatVisitor visitor = new MySqlSchemaStatVisitor();
        sqlStatements.get(i).accept(visitor);
        Map<TableStat.Name, TableStat> tables = visitor.getTables();
        Set<TableStat.Name> names = tables.keySet();
        String table = "";
        for (TableStat.Name name : names) {
            table = name.toString();
        }
        //如果没带库名则通过USE 语句获取库名
        if (!table.contains(".")) {
            String database;
            //获取库名
            List<SQLStatement> list = new ArrayList();
            for (int j = 0; j <= i; j++) {
                list.add(sqlStatements.get(j));
            }
            Collections.reverse(list);
            for (SQLStatement sqlStatement : list) {
                if (sqlStatement.toString().startsWith(SqlCheckUtil.USE)) {
                    String use = sqlStatement.toString();
                    if (!use.endsWith(";")) {
                        throw new RuntimeException("结尾没有分号:" + use);
                    }
                    database = use.substring(use.indexOf(" "), use.indexOf(";"));
                    table = database + "." + table;
                    break;
                }
            }
        }
        table = table.replaceAll(" ", "");
        return table.replaceAll("`", "");
    }

    public static String getSqlType(List<SQLStatement> sqlStatements) {
        for (SQLStatement sqlStatement : sqlStatements) {
            if (sqlStatement instanceof SQLAlterStatement) {
                return DDL;
            } else if (sqlStatement instanceof SQLCreateStatement || sqlStatement instanceof SQLInsertStatement || sqlStatement instanceof SQLDeleteStatement || sqlStatement instanceof SQLUpdateStatement || sqlStatement instanceof SQLReplaceStatement) {
                return DML;
            }
        }
        return USE;
    }

    public static String checkCreateGrammar(List<SQLStatement> sqlStatements) {
        for (SQLStatement sqlStatement : sqlStatements) {
            if (sqlStatement instanceof SQLCreateTableStatement) {
                SQLCreateTableStatement sqlCreateTableStatement = (SQLCreateTableStatement) sqlStatement;
                if (sqlCreateTableStatement.getLike() != null) {
                    continue;
                }
                if (sqlCreateTableStatement.getSelect() != null) {
                    throw new RuntimeException("create table xx select 子句不会创建索引，并且这种语句在复制是不安全的");
                }
                boolean charsetFlag = false;
                if(sqlCreateTableStatement.toString().toUpperCase().contains(SqlCheckUtil.CHARSET)
                        || sqlCreateTableStatement.toString().toUpperCase().contains(SqlCheckUtil.CHARACTER_SET)) {
                    charsetFlag = true;
                }
                if (!charsetFlag) {
                    throw new RuntimeException("建表语句请指定字符集\n" + sqlStatement.toString());
                }
            }
        }
        return null;
    }

    public static void checkAlterGrammar(List<SQLStatement> sqlStatements) {
        Map<String, Integer> alterTable = new HashMap<>();
        for (int i = 0; i < sqlStatements.size(); i++) {
            SQLStatement sqlStatement = sqlStatements.get(i);
            if (sqlStatement instanceof SQLAlterTableStatement) {
                String tableName = getTableName(sqlStatements, i);
                Integer orDefault = alterTable.getOrDefault(tableName, 0);
                alterTable.put(tableName, ++orDefault);
            } else if (sqlStatement instanceof MySqlRenameTableStatement) {
                throw new RuntimeException("暂不支持修改表名,请联系DBA处理");
            }
        }
        for (Map.Entry<String, Integer> entry : alterTable.entrySet()) {
            Integer value = entry.getValue();
            if (value > 1) {
                String key = entry.getKey();
                throw new RuntimeException("合并表`" + key + "`的alter语句以减少操作次数");
            }
        }
    }

    public static List<SQLStatement> parseStatements(String command) {
        try {
            return SQLUtils.parseStatements(command, DbConstants.MYSQL.toLowerCase());
        } catch (Exception e) {
            throw new RuntimeException("请输入有效SQL：" + e.getMessage());
        }
    }

    //添加limit限制,要求传入是检测的sql
    public static String addLimit(String sql, int currentPage, int pageSize) {

        sql = sql.trim();
        if (sql.endsWith(";")) {
            sql = sql.substring(0, sql.length() - 1);
        }
        SQLStatementParser sqlStatementParser = SQLParserUtils.createSQLStatementParser(sql, DbConstants.MYSQL.toLowerCase());
        SQLStatement sqlStatement = sqlStatementParser.parseSelect();
        if (sqlStatement instanceof SQLSelectStatement) {
            SQLSelectStatement queryBlock = (SQLSelectStatement) sqlStatement;
            SQLSelect select = queryBlock.getSelect();
            SQLSelectQueryBlock selectQueryBlock = select.getQueryBlock();
            int noOffsetNum = (currentPage - 1) * pageSize;
            if (selectQueryBlock != null) {
                SQLLimit limit = selectQueryBlock.getLimit();

                if (limit != null) {
                    updateLimit(pageSize, limit, noOffsetNum);
                    return selectQueryBlock.toString();
                } else {
                    sql = selectQueryBlock.toString() + " \nlimit " + noOffsetNum + "," + pageSize;
                }
            } else {
                SQLSelectQuery query = select.getQuery();
                SQLUnionQuery sqlUnionQuery = (SQLUnionQuery) query;
                MySqlSelectQueryBlock rightUnionQuery = getRightUnionQuery(sqlUnionQuery);
                SQLLimit limit = rightUnionQuery.getLimit();
                if (limit != null) {
                    updateLimit(pageSize, limit, noOffsetNum);
                    return sqlUnionQuery.toString();

                } else {
                    sql = sqlUnionQuery.toString() + " \nlimit " + noOffsetNum + "," + pageSize;
                }
            }
        }
        return sql;
    }

    private static MySqlSelectQueryBlock getRightUnionQuery(SQLUnionQuery sqlUnionQuery ){
        SQLSelectQuery right = sqlUnionQuery.getRight();
        if (right instanceof MySqlSelectQueryBlock){
            return ((MySqlSelectQueryBlock) right);
        }else if (right instanceof SQLUnionQuery){
            return getRightUnionQuery(((SQLUnionQuery) right));
        }
        return null;
    }

    private static void updateLimit(int pageSize, SQLLimit limit, int noOffsetNum) {
        SQLExpr rowCount = limit.getRowCount();
        //sql中的偏移量
        int limitOffset = 0;
        if (limit.getOffset() != null) {
            limitOffset = Integer.parseInt(limit.getOffset().toString());
        }
        //分页后的偏移量
        int offset = limitOffset + noOffsetNum;
        //查询总行数
        int totalSize = Integer.parseInt(rowCount.toString());
        int realPageSize = totalSize - noOffsetNum;
        if (realPageSize <= pageSize && realPageSize > 0) {
            limit.setRowCount(realPageSize);
        } else if (realPageSize > pageSize) {
            limit.setRowCount(pageSize);
        } else {
            throw new RuntimeException("超出SQL中标识的查询总数");
        }
        limit.setOffset(offset);
    }
}
